// c-api-examples/asr-with-speaker-identification-c-api.c
//
// Copyright (c)  2024  Xiaomi Corporation

// This file shows how to use sherpa-onnx C API
// to perform ASR with speaker identification.

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <dirent.h>


#include "sherpa-onnx/c-api/c-api.h"

// Structure to hold speaker information
typedef struct {
    const char* name;
    const float* embedding;
} SpeakerInfo;

// Function to compute speaker embedding
static const float* ComputeEmbedding(
    const SherpaOnnxSpeakerEmbeddingExtractor* ex, const char* wav_filename) {
    const SherpaOnnxWave* wave = SherpaOnnxReadWave(wav_filename);
    if (wave == NULL) {
        fprintf(stderr, "Failed to read %s\n", wav_filename);
        exit(-1);
    }

    const SherpaOnnxOnlineStream* stream =
        SherpaOnnxSpeakerEmbeddingExtractorCreateStream(ex);

    SherpaOnnxOnlineStreamAcceptWaveform(stream, wave->sample_rate, wave->samples,
                                       wave->num_samples);
    SherpaOnnxOnlineStreamInputFinished(stream);

    if (!SherpaOnnxSpeakerEmbeddingExtractorIsReady(ex, stream)) {
        fprintf(stderr, "The input wave file %s is too short!\n", wav_filename);
        exit(-1);
    }

    const float* v =
        SherpaOnnxSpeakerEmbeddingExtractorComputeEmbedding(ex, stream);

    SherpaOnnxDestroyOnlineStream(stream);
    SherpaOnnxFreeWave(wave);

    return v;
}

// Function to register speakers from a directory
static int32_t RegisterSpeakers(
    const SherpaOnnxSpeakerEmbeddingManager* manager,
    const SherpaOnnxSpeakerEmbeddingExtractor* ex,
    const char* speakers_dir) {
    
    DIR* dir = opendir(speakers_dir);
    if (!dir) {
        fprintf(stderr, "Failed to open directory: %s\n", speakers_dir);
        return -1;
    }

    struct dirent* entry;
    while ((entry = readdir(dir)) != NULL) {
        // Skip . and .. directories
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0) {
            continue;
        }

        // Check if it's a directory (each speaker should have their own directory)
        char speaker_path[1024];
        snprintf(speaker_path, sizeof(speaker_path), "%s/%s", speakers_dir, entry->d_name);
        
        DIR* speaker_dir = opendir(speaker_path);
        if (!speaker_dir) {
            continue;
        }

        // Process all wav files in the speaker's directory
        struct dirent* wav_entry;
        const float* embeddings[10] = {NULL};  // Support up to 10 samples per speaker
        int32_t num_samples = 0;

        while ((wav_entry = readdir(speaker_dir)) != NULL && num_samples < 10) {
            if (strstr(wav_entry->d_name, ".wav") == NULL) {
                continue;
            }

            char wav_path[1024];
            snprintf(wav_path, sizeof(wav_path), "%s/%s", speaker_path, wav_entry->d_name);
            
            embeddings[num_samples] = ComputeEmbedding(ex, wav_path);
            if (embeddings[num_samples]) {
                num_samples++;
            }
        }

        closedir(speaker_dir);

        // Register speaker if we have samples
        if (num_samples > 0) {
            if (!SherpaOnnxSpeakerEmbeddingManagerAddList(manager, entry->d_name, embeddings)) {
                fprintf(stderr, "Failed to register speaker: %s\n", entry->d_name);
            } else {
                fprintf(stderr, "Successfully registered speaker: %s with %d samples\n", 
                        entry->d_name, num_samples);
            }

            // Free embeddings
            for (int32_t i = 0; i < num_samples; i++) {
                if (embeddings[i]) {
                    SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(embeddings[i]);
                }
            }
        }
    }

    closedir(dir);
    return 0;
}

// Function to identify speaker from audio segment
static const char* IdentifySpeaker(
    const SherpaOnnxSpeakerEmbeddingManager* manager,
    const SherpaOnnxSpeakerEmbeddingExtractor* ex,
    const float* samples,
    int32_t n_samples,
    int32_t sample_rate,
    float threshold) {
    
    const SherpaOnnxOnlineStream* stream =
        SherpaOnnxSpeakerEmbeddingExtractorCreateStream(ex);

    SherpaOnnxOnlineStreamAcceptWaveform(stream, sample_rate, samples, n_samples);
    SherpaOnnxOnlineStreamInputFinished(stream);

    if (!SherpaOnnxSpeakerEmbeddingExtractorIsReady(ex, stream)) {
        SherpaOnnxDestroyOnlineStream(stream);
        return NULL;
    }

    const float* embedding = SherpaOnnxSpeakerEmbeddingExtractorComputeEmbedding(ex, stream);
    const char* speaker = SherpaOnnxSpeakerEmbeddingManagerSearch(manager, embedding, threshold);
    
    SherpaOnnxSpeakerEmbeddingExtractorDestroyEmbedding(embedding);
    SherpaOnnxDestroyOnlineStream(stream);
    
    return speaker;
}

int32_t main(int32_t argc, char* argv[]) {
    if (argc < 3) {
        fprintf(stderr, "Usage: %s <wav_file> <speakers_dir> [speaker_embedding_model] [asr_model_dir]\n", argv[0]);
        fprintf(stderr, "\nSpeakers directory structure should be:\n");
        fprintf(stderr, "speakers_dir/\n");
        fprintf(stderr, "  ├── speaker1/\n");
        fprintf(stderr, "  │   ├── sample1.wav\n");
        fprintf(stderr, "  │   ├── sample2.wav\n");
        fprintf(stderr, "  │   └── ...\n");
        fprintf(stderr, "  ├── speaker2/\n");
        fprintf(stderr, "  │   ├── sample1.wav\n");
        fprintf(stderr, "  │   └── ...\n");
        fprintf(stderr, "  └── ...\n");
        return -1;
    }

    const char* wav_filename = argv[1];
    const char* speakers_dir = argv[2];
    const char* speaker_model = argc > 3 ? argv[3] : "./3dspeaker_speech_campplus_sv_zh-cn_16k-common.onnx";
    const char* asr_model_dir = argc > 4 ? argv[4] : "./sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20";

    // Initialize speaker embedding extractor
    SherpaOnnxSpeakerEmbeddingExtractorConfig speaker_config;
    memset(&speaker_config, 0, sizeof(speaker_config));
    speaker_config.model = speaker_model;
    speaker_config.num_threads = 1;
    speaker_config.debug = 0;
    speaker_config.provider = "cpu";

    const SherpaOnnxSpeakerEmbeddingExtractor* speaker_ex =
        SherpaOnnxCreateSpeakerEmbeddingExtractor(&speaker_config);
    if (!speaker_ex) {
        fprintf(stderr, "Failed to create speaker embedding extractor\n");
        return -1;
    }

    // Initialize speaker embedding manager
    int32_t dim = SherpaOnnxSpeakerEmbeddingExtractorDim(speaker_ex);
    const SherpaOnnxSpeakerEmbeddingManager* speaker_manager =
        SherpaOnnxCreateSpeakerEmbeddingManager(dim);

    // Register speakers
    if (RegisterSpeakers(speaker_manager, speaker_ex, speakers_dir) != 0) {
        fprintf(stderr, "Failed to register speakers\n");
        return -1;
    }

    // Initialize ASR recognizer
    SherpaOnnxOnlineRecognizerConfig asr_config;
    memset(&asr_config, 0, sizeof(asr_config));
    
    asr_config.model_config.tokens = strcat(strdup(asr_model_dir), "/tokens.txt");
    asr_config.model_config.transducer.encoder = strcat(strdup(asr_model_dir), "/encoder.onnx");
    asr_config.model_config.transducer.decoder = strcat(strdup(asr_model_dir), "/decoder.onnx");
    asr_config.model_config.transducer.joiner = strcat(strdup(asr_model_dir), "/joiner.onnx");
    
    asr_config.model_config.num_threads = 1;
    asr_config.model_config.debug = 0;
    asr_config.model_config.provider = "cpu";

    asr_config.decoding_method = "greedy_search";
    asr_config.max_active_paths = 4;

    asr_config.feat_config.sample_rate = 16000;
    asr_config.feat_config.feature_dim = 80;

    asr_config.enable_endpoint = 1;
    asr_config.rule1_min_trailing_silence = 2.4;
    asr_config.rule2_min_trailing_silence = 1.2;
    asr_config.rule3_min_utterance_length = 300;

    const SherpaOnnxOnlineRecognizer* recognizer =
        SherpaOnnxCreateOnlineRecognizer(&asr_config);
    if (!recognizer) {
        fprintf(stderr, "Failed to create recognizer\n");
        return -1;
    }

    // Read wave file
    const SherpaOnnxWave* wave = SherpaOnnxReadWave(wav_filename);
    if (wave == NULL) {
        fprintf(stderr, "Failed to read %s\n", wav_filename);
        return -1;
    }

    // Create stream
    const SherpaOnnxOnlineStream* stream = SherpaOnnxCreateOnlineStream(recognizer);
    const SherpaOnnxDisplay* display = SherpaOnnxCreateDisplay(50);
    int32_t segment_id = 0;

    // Process audio in chunks
    #define N 3200  // 0.2 s. Sample rate is fixed to 16 kHz
    int32_t k = 0;
    while (k < wave->num_samples) {
        int32_t start = k;
        int32_t end = (start + N > wave->num_samples) ? wave->num_samples : (start + N);
        k += N;

        // Process ASR
        SherpaOnnxOnlineStreamAcceptWaveform(stream, wave->sample_rate,
                                           wave->samples + start, end - start);
        while (SherpaOnnxIsOnlineStreamReady(recognizer, stream)) {
            SherpaOnnxDecodeOnlineStream(recognizer, stream);
        }

        const SherpaOnnxOnlineRecognizerResult* r =
            SherpaOnnxGetOnlineStreamResult(recognizer, stream);

        if (strlen(r->text)) {
            // Identify speaker for this segment
            const char* speaker = IdentifySpeaker(speaker_manager, speaker_ex,
                                                wave->samples + start, end - start,
                                                wave->sample_rate, 0.6);
            
            if (speaker) {
                fprintf(stderr, "[%s] ", speaker);
                SherpaOnnxSpeakerEmbeddingManagerFreeSearch(speaker);
            }
            
            SherpaOnnxPrint(display, segment_id, r->text);
        }

        if (SherpaOnnxOnlineStreamIsEndpoint(recognizer, stream)) {
            if (strlen(r->text)) {
                ++segment_id;
            }
            SherpaOnnxOnlineStreamReset(recognizer, stream);
        }

        SherpaOnnxDestroyOnlineRecognizerResult(r);
    }

    // Cleanup
    SherpaOnnxFreeWave(wave);
    SherpaOnnxDestroyDisplay(display);
    SherpaOnnxDestroyOnlineStream(stream);
    SherpaOnnxDestroyOnlineRecognizer(recognizer);
    SherpaOnnxDestroySpeakerEmbeddingManager(speaker_manager);
    SherpaOnnxDestroySpeakerEmbeddingExtractor(speaker_ex);

    return 0;
} 